from tokencal import calculate_tokens
import re

role_info = {
    "finder": "You are a professional academic paper finder.",
    "writer": "You are a professional academic paper writer.",
    "reviewer": "You are a professional academic paper reviewer.",
}

class Prompt:
    def __init__(self, template: str, can_replace: list[str]):
        self.template = template
        self.can_replace = can_replace
        self.replace_dict = {}
        for c in can_replace:
            self.replace_dict[c] = ""

    def what_can_replace(self):
        return self.can_replace
    
    def do(self, placeholder: str, content) -> 'Prompt':
        if placeholder not in self.can_replace:
            raise ValueError(f"Placeholder {placeholder} is not in the can_replace list")
        self.replace_dict[placeholder] = content
        return self.output
    
    @property
    def tokens(self):
        return calculate_tokens(self.output)
    
    @property
    def output(self):
        return self.template.format(**self.replace_dict)
    
    def __str__(self):
        return self.output

problem_classification = Prompt(
"""
1 对一个技术概念的调研综述
2 对一个方向研究现状的综述
3 对多篇文献的对比分析综述
4 对一个技术方法研究脉络的综述
你是一个问题分类专家，请你将用户的输入分类到上面的类别中，仅输出一个数字表示类别。

输入：损失函数，输出：1
输入：Text2SQL研究现状如何，面临哪些挑战，输出：2
输入：有哪些方法可以提升大模型的规划能力，各自优劣是什么，输出：3
输入：多模态大模型的技术发展路线是什么样的，输出：4
输入：{problem}，输出：
""",
    ["problem"]
)

keyword_expand = Prompt(
"""
你是一个科研人员，需要根据某个研究内容搜集文献，请你联想和这个研究主题强相关的{num}个英文短语或句子用于检索相关文献，输出每行一个仅输出短语。/
主题：{theme}
输出：
""",
    ["num", "theme"]
)

description_expand = Prompt(
"""
You are a researcher. Please write a short passage about the historical development of a research topic and the related technologies.
Topic: {theme}
Output:
""",
    ["theme"]
)

keyword_expand_from_description = Prompt(
"""
You are a researcher. Please extract keywords from the theme and description for searching related papers.
The keywords should be related to the theme or the direct technologies.
For example, "Initial stage (1990s)" and "Development history" are invalid keywords.
Theme: {theme}
Description: {description}
Output:
""",
    ["theme", "description"]
)

keyword_expand_from_section = Prompt(
"""
你是一个科研人员，为了写一份综述你需要搜集文献，现在给你主题和综述的一个部分的内容描述，从中抽取英文短语或句子用于检索相关文献，输出每行一个仅输出短语。
主题：{theme}
内容描述：{description}
输出：
""",
    ["theme", "description"]
)

keyword_relativity_score = Prompt(
    """
    You are a expert in academic research. Please score the relevance of a keyword to a research topic.
    The result should be one of the following: "strongly", "weakly", "not".
    Output only the result without any other information.
    Theme: {theme}
    Keyword: {keyword}
    Result:
    """,
    ["theme", "keyword"]
)

relativity_score = Prompt(
"""
You are a expert in academic research. Please score the relevance of a paper fragment to a research topic.
The result should be one of the following: "strongly", "weakly", "not".
Output only the result without any other information.
Theme: {theme}
Paper fragment: {paper_fragment}
Result:
""",
    ["theme", "paper_fragment"]
)

summary_weakly_chunk = Prompt(
"""
You are a expert in academic research. Please summarize the paper fragment.
Output only the result without any other information.
Paper fragment: {paper_fragment}
Result:
""",
    ["paper_fragment"]
)

rough_outline = Prompt(
"""
You wants to write a overall and comprehensive academic survey about "{topic}".
You are provided with a list of papers related to the topic below:
---
{paper_list}
---
You need to draft a outline based on the given papers.
The outline should contains a title and several sections.
Each section follows with a brief sentence to describe what to write in this section.
The outline is supposed to be comprehensive and contains {section_num} sections.

Return in the format:
<format>
Title: A Survey on {topic}
Section 1: [NAME OF SECTION 1]
Description 1: [DESCRIPTION OF SENTCTION 1]

Section 2: [NAME OF SECTION 2]
Description 2: [DESCRIPTION OF SENTCTION 2]

...

Section K: [NAME OF SECTION K]
Description K: [DESCRIPTION OF SENTCTION K]
</format>
The outline:
""",
    ["paper_list", "section_num", "topic"]
)

def rough_outline_parser(outline):
    title = outline.split('Title: ')[1].split('\n')[0]
    sections, descriptions = [], []
    i = 1
    while f'Section {i}' in outline:
        sections.append(outline.split(f'Section {i}: ')[1].split('\n')[0])
        descriptions.append(outline.split(f'Description {i}: ')[1].split('\n')[0])
        i += 1
    return title, sections, descriptions

sub_section_outline = Prompt(
'''
You are an expert in artificial intelligence who wants to write a overall survey about {topic}.
You have created a overall outline below:
---
{overall_outline}
---
The outline contains a title and several sections.
Each section follows with a brief sentence to describe what to write in this section.
<instruction>
You need to enrich the section {section_name}.
The description of {section_name}: {section_description}
You need to generate the framwork containing several subsections based on the overall outlines.
Each subsection follows with a brief sentence to describe what to write in this subsection.
Don't create too many sub-sections, and avoid covering content that belongs to other parts of the main chapter. For example, if the main chapter includes sections like "Case Studies and Applications" or "Challenges and Future Directions" later on, there's no need to include them in your part.
Look carefully at the papers provided for references:
---
{paper_list}
---
Return the outline in the format:
<format>
Subsection 1: [NAME OF SUBSECTION 1]
Description 1: [DESCRIPTION OF SUBSENTCTION 1]

Subsection 2: [NAME OF SUBSECTION 2]
Description 2: [DESCRIPTION OF SUBSENTCTION 2]

...

Subsection K: [NAME OF SUBSECTION K]
Description K: [DESCRIPTION OF SUBSENTCTION K]
</format>
</instruction>
Only return the outline without any other informations:
''',
    ["overall_outline", "section_name", "section_description", "paper_list", "topic"]
)

def subsections_parser(outline):
    if not outline:
        return [], []

    # 使用正则表达式匹配所有的Subsection和Description
    subsection_pattern = re.compile(r'Subsection (\d+): (.*?)\n')
    description_pattern = re.compile(r'Description (\d+): (.*?)\n')

    subsections = subsection_pattern.findall(outline)
    descriptions = description_pattern.findall(outline)

    # 将匹配结果转换为字典以便按ID查找
    subsection_dict = {int(id_): title for id_, title in subsections}
    description_dict = {int(id_): desc for id_, desc in descriptions}

    # 获取所有Subsection的ID并按顺序处理
    subsection_ids = sorted(subsection_dict.keys())
    subsections_list = []
    descriptions_list = []

    for id_ in subsection_ids:
        subsections_list.append(subsection_dict[id_])
        # 如果对应的Description存在，则添加；否则添加空字符串
        descriptions_list.append(description_dict.get(id_, ""))

    return subsections_list, descriptions_list

advice_on_modify_outline = Prompt(
"""
You wants to write a overall and comprehensive academic survey about "{topic}".
You have created a overall outline below:
---
{overall_outline}
---
The outline contains a title and several sections.
Each section follows with a brief sentence to describe what to write in this section.
Under each section, there are several subsections.
Each subsection also follows with a brief sentence of descripition.
Some of the subsections may be repeated or overlaped.
I need to make sure the outline is both comprehensive and logically coherent with no repeated subsections.
Repeated subsections among sections are not allowed! Also, make sure that the sub-sections do not conflict with the main chapter, for example, if the main chapter includes sections like "Case Studies and Applications" or "Challenges and Future Directions" later on, there's no need to include them in this part.
I want you to give me some specific advice based on tips above to modify the outline. But don't modify the outline, just tell me what to modify.
""",
    ["overall_outline", "topic"]
)

modify_outline = Prompt(
"""
You wants to write a overall and comprehensive academic survey about "{topic}".
You have created a overall outline below:
---
{overall_outline}
---
You have received some advice below:
---
{advice}
---
You need to modify the outline based on the advice.
Return the final outline in the format:
<format>
Title: {topic}
Section 1: [NAME OF SECTION 1]
Description 1: [DESCRIPTION OF SENTCTION 1]

Subsection 1: [NAME OF SUBSECTION 1]
Description 1: [DESCRIPTION OF SUBSENTCTION 1]

Subsection 2: [NAME OF SUBSECTION 2]
Description 2: [DESCRIPTION OF SUBSENTCTION 2]

...

Section K: [NAME OF SECTION K]
Description K: [DESCRIPTION OF SENTCTION K]
</format>
The outline:
""",
    ["overall_outline", "topic", "advice"]
)

def modify_outline_parser(outline):
    # 解析标题
    title_match = re.search(r"^Title:\s*(.+)", outline, re.MULTILINE)
    title = title_match.group(1) if title_match else "Untitled"

    # 解析章节
    section_pattern = re.compile(r"Section (\d+): (.+)")
    subsection_pattern = re.compile(r"Subsection (\d+): (.+)")
    description_pattern = re.compile(r"Description (\d+): (.+)")

    sections = []
    current_section = None
    current_subsections = []

    for line in outline.split("\n"):
        section_match = section_pattern.match(line)
        subsection_match = subsection_pattern.match(line)
        description_match = description_pattern.match(line)

        if section_match:
            if current_section:
                current_section["subsections"] = current_subsections
                sections.append(current_section)
            current_section = {
                "name": section_match.group(2),
                "description": "",
                "subsections": []
            }
            current_subsections = []
        
        elif subsection_match:
            current_subsections.append({
                "subsection_name": subsection_match.group(2),
                "subsection_description": ""
            })
        
        elif description_match:
            desc_text = description_match.group(2)
            if current_subsections and int(description_match.group(1)) == len(current_subsections):
                current_subsections[-1]["subsection_description"] = desc_text
            elif current_section and not current_section["description"]:
                current_section["description"] = desc_text

    if current_section:
        current_section["subsections"] = current_subsections
        sections.append(current_section)

    result = {
        "title": title,
        "sections": sections
    }
    return result

arrange_papers = Prompt(
"""
You are an expert in academic research, assisting in writing a literature review.  
Your task is to categorize the given paper fragment into the most appropriate subsection based on the provided outline.  

---
**Paper**  
- **Title**: {paper_title}  
- **Content**: {paper_text}  

---
**Outline**  
{overall_outline}  

---
Return only the subsection(s) where the paper fragment belongs, formatted as a JSON list (e.g., ["1.1", "2.6"]). 
If none of the subsections are suitable, return an empty list.
""",
    ["paper_title", "paper_text", "overall_outline"]
)

def arrange_papers_parser(result):
    import json
    if "```" in result:
        result = result.strip("```").strip("json").strip()
    try:
        return json.loads(result)
    except:
        print(f"Error parsing JSON result: {result}")
        return []


find_table_insert_position = Prompt(
"""
You are an expert in academic research. 
Please determine whether a table is needed to be added to this section to compare methods or illustrate the progress of research (If the references are too few or the section is too short you should choose "No"). Please consider carefully and choose "yes" only if it is truly meaningful. Answer with only "Yes" or "No".
---
Section:
{section}
---
Result:
""",
    ["section"]
)

insert_table_prompt = Prompt(
"""
You are an expert in academic research. 
Please insert a table to this section to compare methods or illustrate the progress of research.
The table should include the year(in year order), method (or technology or development stage), a brief description, and the title of the references, with the paper_title and chunk_id enclosed in [].
The result should only include the table in markdown format, without any other information.
For example:
| Year | Method/Technology/Development Stage | Description | References |
|------|-------------------------------------|-------------|------------|
| 2020 | Deep Learning-based Models         | Introduction of deep learning models to automate literature review generation | [Deep Learning for Literature Review Generation, 1] |
| 2021 | Knowledge Graph Construction        | Use of knowledge graphs to improve the semantic understanding of academic papers | [ Building Knowledge Graphs for Research, 2] |
| 2022 | Reinforcement Learning for Review Writing | Application of reinforcement learning to enhance automated literature review writing | [Reinforcement Learning for Writing Reviews, 3] |
| 2023 | Multi-Agent Systems for Collaborative Review | Development of multi-agent systems involving Fetcher, Writer, and Editor roles to improve collaborative review generation | [Multi-Agent Systems for Academic Review, 4] |
---
Section:
{section}
---
Result:
""",
    ["section"]
)

def insert_table_parser(result):
    if "```markdown" in result:
        result = result.strip("```markdown").strip("```")
    return result

subsection_writing = Prompt(
'''
You are an expert who wants to write a overall and comprehensive survey about {topic}.
You have created a overall outline below:
---
{overall_outline}
---
Below are a list of papers for references:
---
{paper_list}
---

<instruction>
Now you need to write the content for the subsection:
"{subsection_name}" under the section: "{section_name}"
The details of what to write in this subsection called "{subsection_name}" is in this descripition:
---
{description}
---

Here is the requirement you must follow:
1. When writing sentences that are based on specific papers above, you cite the "paper_title" and "chunk_id" in a '[]' format to support your content. An example of citation: 'the emergence of large language models (LLMs) [Language models are few-shot learners, 2]'
Note that the "paper_title" is not allowed to appear without a '[]' format. Also remember to keep the full paper title in '[]'. Papers not existing above are not allowed to cite!
2. When cite the content in the paper, don't directly copy the paper title outside [], use the method name or the main idea of the paper to describe the content.
3. Don't need to use words like "firstly" or "in conclusion" in your writing, as you are only responsible for a part of the content, and it will be connected to other sections before and after. Please keep this in mind.
4. Make full use of the provided "paper_list". Keep the content concise and to the point, and avoid being boring and repetitive.

Here's a concise guideline for when to cite papers in a survey:
---
1. Summarizing Research: Cite sources when summarizing the existing literature.
2. Using Specific Concepts or Data: Provide citations when discussing specific theories, models, or data.
3. Comparing Findings: Cite relevant studies when comparing or contrasting different findings.
4. Highlighting Research Gaps: Cite previous research when pointing out gaps your survey addresses.
5. Using Established Methods: Cite the creators of methodologies you employ in your survey.
6. Supporting Arguments: Cite sources that back up your conclusions and arguments.
7. Suggesting Future Research: Reference studies related to proposed future research directions.
---

</instruction>
Return the content of subsection "{subsection_name}" in the format:
<format>
content
</format>
''',
    ["overall_outline", "subsection_name", "section_name", "description", "paper_list", "topic"]
)

def subsection_writing_parser(result):
    return result.strip("<format>").strip("</format>")

abstract_writing = Prompt(
'''
You are an expert in artificial intelligence who wants to write a overall and comprehensive survey about {topic}.
You have created a overall passage below:
---
{overall_passage}
---
You need to write a abstract for the survey.
The abstract should be concise and to the point.

Return in the format:
<format>
Abstract: 
</format>
Only return the abstract without any other information:
''',
    ["overall_passage", "topic"]
)

def extract_citations(markdown_text):
    # 正则表达式匹配方括号内的内容
    pattern = re.compile(r'\[(.*?)\]')
    matches = pattern.findall(markdown_text)
    # 分割引用，处理多引用情况，并去重
    citations = list()
    for match in matches:
        # 分割各个引用并去除空格
        parts = match.split(';')
        for part in parts:
            cit = part.strip()
            if cit not in citations:
                citations.append(cit)
    return citations

def final_format(final_md:dict):
    final_text = ""
    final_text += f"# {final_md['title']}\n\n"
    abstract = final_md['abstract'].strip("<format>").strip("</format>")
    final_text += f"{abstract}\n\n"
    for section in final_md['sections']:
        final_text += f"## {section['title']}\n\n"
        for subsection in section['subsections']:
            final_text += f"### {subsection['title']}\n\n"
            content = subsection['content'].strip("<format>").strip("</format>")
            final_text += f"{content}\n\n"
    return final_text


if __name__ == "__main__":
    print(problem_classification.do("problem", "多模态大模型的发展历程"))
    print(problem_classification.tokens)